#include "xparameters.h"
#include "xuartps.h"
#include "xil_io.h"
#include "mnist_params.h"
#include <limits.h>

#define UART_DEVICE_ID   XPAR_XUARTPS_0_DEVICE_ID
#define UART_BASEADDR    XPAR_XUARTPS_0_BASEADDR
#define BRAM_BASEADDR    XPAR_BRAM_0_BASEADDR  // 根据 IP 名称可能是 BRAM_CTRL_0_BRAM_PORTA_BASEADDR
#define IMAGE_SIZE       (28*28)

int main(void)
{
    XUartPs uart;
    XUartPs_Config *cfg;
    u8 buf[IMAGE_SIZE];
    int idx = 0;

    // 初始化 UART
    cfg = XUartPs_LookupConfig(UART_DEVICE_ID);
    XUartPs_CfgInitialize(&uart, cfg, UART_BASEADDR);
    XUartPs_SetBaudRate(&uart, 115200);

    while (1) {
        // 等待帧头 “START”
        while (1) {
            if (XUartPs_IsReceiveData(UART_BASEADDR)) {
                char c = XUartPs_RecvByte(UART_BASEADDR);
                if (c == 'S') {
                    // 简单匹配五字节 “START”
                    char tmp[5];
                    tmp[0]=XUartPs_RecvByte(UART_BASEADDR);
                    tmp[1]=XUartPs_RecvByte(UART_BASEADDR);
                    tmp[2]=XUartPs_RecvByte(UART_BASEADDR);
                    tmp[3]=XUartPs_RecvByte(UART_BASEADDR);
                    if (tmp[0]=='T'&&tmp[1]=='A'&&tmp[2]=='R'&&tmp[3]=='T')
                        break;
                }
            }
        }

        // 读 784 字节像素
        for (idx = 0; idx < IMAGE_SIZE; idx++) {
            while (!XUartPs_IsReceiveData(UART_BASEADDR));
            buf[idx] = XUartPs_RecvByte(UART_BASEADDR);
        }

        // 等待帧尾 “END”
        while (1) {
            if (XUartPs_IsReceiveData(UART_BASEADDR)) {
                char c = XUartPs_RecvByte(UART_BASEADDR);
                if (c=='E'
                    && XUartPs_RecvByte(UART_BASEADDR)=='N'
                    && XUartPs_RecvByte(UART_BASEADDR)=='D')
                    break;
            }
        }

        // 写入 BRAM
        for (idx = 0; idx < IMAGE_SIZE; idx++) {
            Xil_Out8(BRAM_BASEADDR + idx, buf[idx]);
        }

        // 回发 DONE 表示数据已经写入
        const char done[] = "DONE";
        XUartPs_Send(&uart, (u8 *)done, sizeof(done)-1);

        // 在 FPGA 端执行 MNIST 推理
        int best = INT_MIN;
        int label = 0;
        for (int k = 0; k < 10; k++) {
            int sum = b[k];
            for (int i = 0; i < IMAGE_SIZE; i++) {
                sum += W[k][i] * buf[i];
            }
            if (sum > best) {
                best = sum;
                label = k;
            }
        }
        // 发送识别结果 ASCII
        char out = '0' + label;
        XUartPs_Send(&uart, (u8 *)&out, 1);

        // 处理完毕后循环回到等待下一次 START
    }
    return 0;
}
